#include "../include/SSS_main.cuh"
#include "../include/SSS_matvec.h"
#include "../include/SSS_AMG.h"
#include "../include/mmio.h"
#include "../include/mmio_highlevel.h"
#include "../include/mmio_utils.h"

//#include <cuda_runtime_api.h>
double dot_comp_time = 0;
double dot_time = 0;

double axpy_comp_time = 0;
double axpy_time = 0;

double axpby_comp_time = 0;
double axpby_time = 0;

double ax_comp_time = 0;
double ax_time = 0;

double rap_time = 0;

double solve_time =0;
double setup_time =0;
double total_time =0;

//mxy
 double mxy_time =0 ;

//amxpy
 double amxpy_time =0;

 int output_matrix = 0 ;
//mxy
void spmv(const SSS_MAT *A, const SSS_VEC *x, SSS_VEC *y)
{
    double d1 = SSS_get_time();
    int m = A->num_rows;
    int *ia = A->row_ptr, *ja = A->col_idx;
    double *Aj = A->val;
    int i, k, begin_row, end_row;
    double temp;

    for (i = 0; i < m; ++i) {
        temp = 0.0;
        begin_row = ia[i];
        end_row = ia[i + 1];

        for (k = begin_row; k < end_row; ++k) {
            temp += Aj[k] * x->d[ja[k]];
        }

        y->d[i] = temp;
    }
    double d2 = SSS_get_time();
    mxy_time+= d2-d1;
}

//读矩阵
void SSS_mat_read(char *filemat, SSS_MAT *A)
{
    
    int isSymmetric;
    printf("filename: %s\n",filemat);
   	mmio_info(&A->num_rows, &A->num_cols, &A->num_nnzs, &isSymmetric, filemat);
    A->row_ptr  = (int *)malloc((A->num_rows+1) * sizeof(int));
	A->col_idx  = (int *)malloc(A->num_nnzs * sizeof(int));
	A->val  = (double *)malloc(A->num_nnzs * sizeof(double));
    mmio_data(A->row_ptr, A->col_idx, A->val, filemat);   
    printf("A: m = %d, n = %d, nnz = %d\n", A->num_rows,A->num_cols, A->num_nnzs);
}


SSS_MAT laplacian_3pt(const int n)
{
    SSS_MAT A;
    int k, left;
    int i;

    assert(n > 1);

    A.num_rows = A.num_cols = n;
    A.num_nnzs = 0;

    left = 64;
    //AP AJ AX
    A.row_ptr = (int*)malloc(sizeof(int) * (n + 1));
    A.col_idx = (int*)malloc(sizeof(int) * left);
    A.val = (double*)malloc(sizeof(double) * left);

    k = 0;
    A.row_ptr[0] = 0;
    for (i = 0; i < A.num_rows; i++) A.row_ptr[i + 1] = 0;

    for (i = 0; i < n; i++) 
    {
        if (left <= 32) {
            left += 64;

            A.col_idx =(int*)realloc(A.col_idx, sizeof(int) * (k + left));
            A.val = (double*)realloc(A.val, sizeof(double) * (k + left));
        }

        if (i == 0) {
            A.col_idx[k] = 0;
            A.val[k] = 2.;
            k++;
            left--;

            A.col_idx[k] = 1;
            A.val[k] = -1;
            k++;
            left--;
        }
        else if (i == n - 1) {
            A.col_idx[k] = n - 2;
            A.val[k] = -1;
            k++;
            left--;

            A.col_idx[k] = n - 1;
            A.val[k] = 2;
            k++;
            left--;
        }
        else {
            A.col_idx[k] = i - 1;
            A.val[k] = -1.;
            k++;
            left--;

            A.col_idx[k] = i;
            A.val[k] = 2.02;
            k++;
            left--;

            A.col_idx[k] = i + 1;
            A.val[k] = -1.;
            k++;
            left--;
        }

        A.row_ptr[i + 1] = k;
    }

    A.num_nnzs = k;

    return A;
}

SSS_MAT laplacian_5pt(const int N)
{
    SSS_MAT A;
    int nz = 0;
    int i, j;

    assert(N > 0);

    A.num_rows = N * N;
    A.num_cols = N * N;
    A.num_nnzs = 5 * N * N;

    A.row_ptr = (int *)malloc(sizeof(int) * (A.num_rows + 1));
    A.col_idx = (int *)malloc(sizeof(int) * A.num_nnzs);
    A.val = (double *)malloc(sizeof(double) * A.num_nnzs);

    A.row_ptr[0] = 0;

    for (i = 0; i < A.num_rows; i++) 
    {
        A.row_ptr[i + 1] = 0;
    }

    for (i = 0; i < N; i++) {
        for (j = 0; j < N; j++) {
            int indx = N * i + j;

            if (i > 0) {
                A.col_idx[nz] = indx - N;
                A.val[nz] = -1;
                nz++;
            }

            if (j > 0) {
                A.col_idx[nz] = indx - 1;
                A.val[nz] = -1;
                nz++;
            }

            A.col_idx[nz] = indx;
            A.val[nz] = 4;
            nz++;

            if (j < N - 1) {
                A.col_idx[nz] = indx + 1;
                A.val[nz] = -1;
                nz++;
            }

            if (i < N - 1) {
                A.col_idx[nz] = indx + N;
                A.val[nz] = -1;
                nz++;
            }

            A.row_ptr[indx + 1] = nz;
        }
    }

    A.num_nnzs = nz;
    return A;
}

SSS_MAT laplacian_7pt_bound(const int Nx, const int Ny, const int Nz)
{
    SSS_MAT A;
    int nz = 0;
    int i, j, k;

    A.num_rows = Nx * Ny * Nz;
    A.num_cols = Nx * Ny * Nz;
    A.num_nnzs = 7 * Nx * Ny * Nz;

    A.row_ptr = (int*)malloc(sizeof(int) * (A.num_rows + 1));
    A.col_idx = (int*)malloc(sizeof(int) * A.num_nnzs);
    A.val = (double*)malloc(sizeof(double) * A.num_nnzs);

    A.row_ptr[0] = 0;
    for (i = 0; i < A.num_rows; i++) A.row_ptr[i + 1] = 0;

    for (i = 0; i < Nz; i++) {
        for (j = 0; j < Ny; j++) {
            for (k = 0; k < Nx; k++) {
                int indx = Nx * Ny * i + Nx * j + k;

                if (i > 0) {
                    A.col_idx[nz] = indx - Nx * Ny;
                    A.val[nz] = -1;
                    nz++;
                }

                if (j > 0) {
                    A.col_idx[nz] = indx - Nx;
                    A.val[nz] = -1;
                    nz++;
                }

                if (k > 0) {
                    A.col_idx[nz] = indx - 1;
                    A.val[nz] = -1;
                    nz++;
                }

                A.col_idx[nz] = indx;
                A.val[nz] = 6;
                nz++;

                if (k < Nx - 1) {
                    A.col_idx[nz] = indx + 1;
                    A.val[nz] = -1;
                    nz++;
                }

                if (j < Ny - 1) {
                    A.col_idx[nz] = indx + Nx;
                    A.val[nz] = -1;
                    nz++;
                }

                if (i < Nz - 1) {
                    A.col_idx[nz] = indx + Nx * Ny;
                    A.val[nz] = -1;
                    nz++;
                }

                A.row_ptr[indx + 1] = nz;
            }
        }
    }

    A.num_nnzs = nz;
    return A;
}

//初始化参数
void SSS_amg_pars_init(SSS_AMG_PARS *pars)
{

    //选择的光滑器
    pars->smoother = SSS_SM_GS;
    //AMG最大迭代次数
    pars->max_it = 100;
    //AMG收敛精度
    pars->tol = 1e-6;
    //最粗层收敛精度
    pars->ctol = 1e-7;
    //AMG最大粗网格层数
    pars->max_levels = 10;
    //最粗层的自由度
    //pars->coarse_dof = MIN_CDOF;
    pars->coarse_dof = 10;
    //循环类型 V-cycle
    pars->cycle_type = 1;
    //cf_order
    //0: nature order      1: C/F order顺序
    pars->cf_order = 1;
    //前光滑次数
    pars->pre_iter = 3;
    //后光滑次数
    pars->post_iter = 3;
    //SOR光滑器的松弛解析器
    // pars->relax = 1.0;
    //多项式平滑度 (Polynomial smoother)
    // pars->poly_deg = 3;
    //粗化类型
    pars->cs_type = SSS_COARSE_RS;
    //插值类型
    pars->interp_type = intERP_DIR;
    //最大行和分析器
    pars->max_row_sum = 0.9;
    //粗化的强连接阈值
    pars->strong_threshold = 0.3;
    //截断阈值（将超过阈值的值设为阈值）
    pars->trunc_threshold = 0.2;
}
 
//打印参数
void SSS_amg_pars_print(SSS_AMG_PARS *pars)
{
    assert(pars != NULL);
    
    printf("\n               AMG Parameters \n");
    printf("-----------------------------------------------------------\n");

    printf("AMG max num of iter:               %d\n", pars->max_it);
    printf("AMG tol:                           %g\n", pars->tol);
    printf("AMG ctol:                          %g\n", pars->ctol);
    printf("AMG max levels:                    %d\n", pars->max_levels);
    printf("AMG cycle type:                    %d\n", pars->cycle_type);
    printf("AMG smoother type:                 %d\n", pars->smoother);
    printf("AMG smoother order:                %d\n", pars->cf_order);
    printf("AMG num of presmoothing:           %d\n", pars->pre_iter);
    printf("AMG num of postsmoothing:          %d\n", pars->post_iter);

    switch(pars->smoother) {
        case SSS_SM_SOR:
        case SSS_SM_SSOR:
        case SSS_SM_GSOR:
        case SSS_SM_SGSOR:
            printf("AMG relax factor:                  %.4lf\n", pars->relax);
            break;

        case SSS_SM_POLY:
            printf("AMG polynomial smoother degree:    %d\n", pars->poly_deg);
            break;

        default:
            break;
    }

    printf("AMG coarsening type:               %d\n", pars->cs_type);
    //printf("AMG interPolation type:            %d\n", pars->interp_type);
    switch (pars->interp_type)
    {
    case 1:
        printf("AMG interPolation type:            Dir\n");
        break;
    case 2:
        printf("AMG interPolation type:            STD\n");
        break;
    default:
        break;
    }
    printf("AMG dof on coarsest grid:          %d\n", pars->coarse_dof);
    printf("AMG strong threshold:              %.4lf\n", pars->strong_threshold);
    printf("AMG truncation threshold:          %.4lf\n", pars->trunc_threshold);
    printf("AMG max row sum:                   %.4lf\n", pars->max_row_sum);

    printf("-----------------------------------------------------------\n");
}

//打印矩阵
void print_matrix(SSS_MAT *A)
{
    int A_nnz = A->num_nnzs;
    int A_row = A->num_rows;
    int A_col = A->num_cols;

    FILE *fp_A;
    char pathA[100];

    char local_time_str[128];
    char *timeA = NULL;

    sprintf(pathA,"Coarsen_matrix/Operator_Matrix_A_m%d_k%d.mtx",A_row,A_col);
    fp_A=fopen(pathA,"a+"); 

    fprintf(fp_A,"%d %d %d\n", A_row, A_col,A_nnz);
    for (int i = 1; i <=A_row; i++)
    {
            for (int j = A->row_ptr[i-1]; j < A->row_ptr[i]; j++)
            {
                //base0(HYPRE)
                //fprintf(fp_A,"%d %d %f\n",i-1,A_j[j],A_data[j]);
                //base1(AMGX)
                fprintf(fp_A,"%d %d %f\n",i,A->col_idx[j]+1,A->val[j]);
            }
    }
    fclose(fp_A);
}
void read_Solver_rhs(char *rhs_name,SSS_VEC *rhs)
{
    FILE* fp;
    fp = fopen(rhs_name, "r");
    if (fp != NULL)
    {
        fscanf(fp,"%d",&rhs->n);
        for(int i=0;i<rhs->n;i++)
        {
            fscanf(fp,"%lf",&rhs->d[i]);
        }
    }
    else
        printf("fail to open! \n");
    fclose(fp);
}

int main(int argc,char *argv[])
{
    int arg_index = 1;

    SSS_AMG_PARS pars;
    SSS_MAT A;
    SSS_VEC b, x;
    SSS_RTN rtn;


    /* 1.读矩阵 处理矩阵 */
    /* read mtx format matrix */
    if ( strcmp(argv[arg_index], "-read_matrix") == 0 )
    {
        arg_index++;
        char *mat_file = argv[arg_index];
        SSS_mat_read(mat_file,&A);
    }
    /* Laplacian 3pt  */
    else if ( strcmp(argv[arg_index], "-laplacian_3pt") == 0 )
    {
        arg_index++;
        A = laplacian_3pt(atoi(argv[arg_index])); 
        printf("3pt Laplacian problem size: %d.\n", atoi(argv[arg_index]));
    }
    /* Laplacian 5pt  */
    else if ( strcmp(argv[arg_index], "-laplacian_5pt") == 0 )
    {
        arg_index++;
        A = laplacian_5pt(atoi(argv[arg_index]));    
        printf("5pt Laplacian problem size: %d, %d x %d.\n", atoi(argv[arg_index]) * atoi(argv[arg_index]), atoi(argv[arg_index]), atoi(argv[arg_index]));
    }
    else if ( strcmp(argv[arg_index], "-laplacian_7pt") == 0 )
    {
        arg_index++;
        int num=atoi(argv[arg_index]);

        A = laplacian_7pt_bound(num,num,num);    
        printf("7pt Laplacian problem size: %d, %d x %d x %d.\n", num*num*num, num, num, num);
    }

    b = SSS_vec_create(A.num_rows);
    int flag_b = 0;

    if (strcmp(argv[arg_index++], "-read_rhs") == 0 )
    {
        arg_index++;
        flag_b = 1;
        char *rhs_file = argv[arg_index];
        printf("rhs_name: %s\n",rhs_file);

        read_Solver_rhs(rhs_file,&b);
    }
    /* print matrix A */
    if(output_matrix==1)
    {
        print_matrix(&A);
    }   
    

    //    for(int i=0;i<A.num_nnzs;i++)
    // {
    //    printf("row_ptr[%d] = %d \n",i,A.row_ptr[i]);
    //     printf("col_idx[%d] = %d \n",i,A.col_idx[i]);
    //     printf("val[%d] = %lf\n",i,A.val[i]);
    // }


    // printf("row_ptr[%d] = %d\n",A.num_rows+1,A.row_ptr[A.num_rows+1]);

    // printf("col_idx[%d] = %d\n",A.num_nnzs,A.col_idx[A.num_nnzs]);

    // printf("val[%d] = %lf\n",A.num_nnzs,A.val[A.num_nnzs]);
    printf("A:m = %d, n = %d, nnz = %d\n",A.num_rows,A.num_cols,A.num_nnzs);
     
    /* pars */
    /* 2.初始化参数 */
    SSS_amg_pars_init(&pars);

    /* print info */
    /* 3.打印AMG参数 */
    SSS_amg_pars_print(&pars);

    /* 4.初始化向量b和x */
    //b = SSS_vec_create(A.num_rows);
    if(flag_b==0)
    {
        SSS_vec_set_value(&b, 1.0);
    }
    


    // for(int i =0;i<10;i++)
    //     printf("b[%d] = %lf\n",i,b.d[i]);
    
    // spmv(&A,&b,&b);
    // for(int i =0;i<10;i++)
    //     printf("bb[%d] = %lf\n",i,b.d[i]);
    

    x = SSS_vec_create(A.num_rows);
    SSS_vec_set_value(&x, 1.0);
    
    /* solve the system */
    /* 5.正式进行AMG */
    rtn = SSS_solver_amg(&A, &x, &b, &pars);
    
    printf("AMG residual: %g\n", rtn.ares);
    printf("AMG relative residual: %g\n", rtn.rres);
    printf("AMG iterations: %d\n", rtn.nits);
    
    // printf("\n");
    // printf("----------------------------------\n");
    // printf("----------------------------------\n");
    // printf("\n");
    // printf("--------------setup_time--------------------\n");

    // printf("setup_time = %lf ms, setup/total = %lf %%\n",setup_time,setup_time/ total_time *100);
    // printf("rap_time   = %lf ms, rap/setup = %lf %% \n",rap_time, rap_time/setup_time *100);

    // printf("\n");
    // printf("--------------solve_time--------------------\n");

    // printf("solve_time = %lf ms, solve/total = %lf %%\n",solve_time,solve_time/ total_time *100);

    SSS_mat_destroy(&A);
    SSS_vec_destroy(&x);
    SSS_vec_destroy(&b);
    return 0;
}